% EKF based localization of a robot with known landmarks

clear; clc; close all;
load Q6_data.mat
%load Q6_example_data.mat

T = length(u);
ss = length(x0);
xfilt = x0;
Sigma_filt{1} = Sigma0;

num_landmarks_discovered = 0;

for t=1:T-1

    A = jacobian_f_robot(xfilt(:,t), u(:,t), dt);
    
    xfilt(:,t+1) = f_robot(xfilt(:,t), u(:,t), dt);
    Sigma_filt{t+1} = A*Sigma_filt{t}*A' + Q*dt;

    landmarks_found_so_far_at_current_time = [];

    for k=1:num_landmark_measurements(t+1)
        y_landmark = landmark_measurements{t+1}(:,k);

        themin_score = Inf;
        lm_score = []; 
    
        % let's see how well the measurement k matches with the
        % hypothesis that it would be a measurement of each landmark
        % we have discovered in the past and current time:
        for l=1:num_landmarks_discovered
            if(~ismember(l, landmarks_found_so_far_at_current_time)) 
                % each landmark can only be observed once, so don't check
                % for landmark already detected in measurements 1 through
                % k-1
                
                C = jacobian_f_unknown_landmark(xfilt(:,t+1), l);
                ypred = f_unknown_landmark(xfilt(:,t+1), l);

                error_y = y_landmark - ypred; % error (innovation)
                S = C*Sigma_filt{t+1}*C' + R_landmark; %Innovation (or residual) covariance

                lm_score(l) = error_y'*pinv(S)*error_y;
            else
                lm_score(l) = Inf; 
            end
            [themin_score, themin_score_lmidx] = min(lm_score);
        end

        if(themin_score < 9)  
            % if within 3 standard deviations of some landmark that we have already discovered
            % then we assume it is that landmark and we perform a standard
            % measurement update:
            
            l = themin_score_lmidx;
            landmarks_found_so_far_at_current_time = [landmarks_found_so_far_at_current_time l];

            C = jacobian_f_unknown_landmark(xfilt(:,t+1), l);
            ypred = f_unknown_landmark(xfilt(:,t+1), l);
            y = y_landmark;
            
            % your code here [perform measurement update]
			K = Sigma_filt{t+1}*C'*inv(C*Sigma_filt{t+1}*C' +  R_landmark);
			xfilt(:,t+1) = xfilt(:,t+1) + K * (y - ypred);
			Sigma_filt{t+1} = Sigma_filt{t+1} - K*C*Sigma_filt{t+1};

        else
            % we found a new landmark not in our state yet:
            % add into filter state:

            % bookkeeping on the landmarks discovered: [use our code so
            % plotting code works later]
            num_landmarks_discovered = num_landmarks_discovered + 1;
            landmark_discovery_time(num_landmarks_discovered) = t + 1;
            landmarks_found_so_far_at_current_time = [landmarks_found_so_far_at_current_time num_landmarks_discovered];

            
            % statesize increases by 2 (the two coordinates of the
            % landmark
            
            ss = ss + 2;

            % augment the state
            % padd zeros for past times for the landmark
            xfilt = [xfilt; zeros(2,size(xfilt,2))];
            
            % compute estimate of landmark position:
            R = [cos(xfilt(3,t+1)) -sin(xfilt(3,t+1)); sin(xfilt(3,t+1)) cos(xfilt(3,t+1))];
			landmark_ne_estimate = R*y_landmark + [xfilt(1,t+1) ; xfilt(2,t+1)] ;
			
            xfilt(end-1:end,t+1) = landmark_ne_estimate;
            
            % augment process noise matrix (no noise on
            % landmarks---they are stationary)
            Q = [ Q  zeros(size(Q,1),2) ; zeros(2, size(Q,2)+2)]; % no process noise on the landmarks

            % augment filter covariance:
			Sigma_filt{t+1} = [Sigma_filt{t+1} zeros(size(Sigma_filt{t+1},1),2); zeros(2, size(Sigma_filt{t+1},2)) zeros(2,2)];
			C = jacobian_f_unknown_landmark(xfilt(:,t+1), num_landmarks_discovered);
			landmark_ne_cov = R' * (C*Sigma_filt{t+1}*C' + 6*R_landmark) * R; %Innovation (or residual) covariance
			Sigma_filt{t+1}(end-1:end,end-1:end) = landmark_ne_cov;

            % [see Q5]

        end
    end
end

colors1 = ['km']; colors2 = ['bg'];
map_fig_id = figure; hold on; axis([-5 25 -5 25]); axis equal; xlabel('East'); ylabel('North');
spacing = 20;
for i= [1:spacing:size(xfilt,2) size(xfilt,2)]
    plot_uncertainty_ellipse(xfilt(1:3,i), Sigma_filt{i}, map_fig_id, colors1);
    for l=1:num_landmarks_discovered
       l_state_idxs = 3+ 2*(l-1)+1: 3 + 2*(l-1)+2;
       if(landmark_discovery_time(l) < i)
           % if we have seen the landmark a few times, start plotting it
           plot_uncertainty_ellipse(xfilt(l_state_idxs,i), Sigma_filt{i}(l_state_idxs,l_state_idxs), map_fig_id, colors2);
       end
    end
end

% plot final map
final_map_fig_id = figure; hold on; axis([-5 25 -5 25]); axis equal; xlabel('East'); ylabel('North'); title('map');
colors3 = ['kr'];
for l=1:num_landmarks_discovered
    l_state_idxs = 3+ 2*(l-1)+1: 3 + 2*(l-1)+2;
    %if(trace(Sigma_filt(l_state_idxs,l_state_idxs,i)) < 2*sigma0_landmarks^2 / 10)
        %% if we have seen the landmark a few times the variance will have decreased and we plot it
        plot_uncertainty_ellipse(xfilt(l_state_idxs,end), Sigma_filt{end}(l_state_idxs,l_state_idxs), final_map_fig_id, colors3);
    %end
end
axis equal;



